import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageDraw, ImageFilter, ImageOps
import numpy as np
import tensorflow as tf

# 加载预训练模型
model = tf.keras.models.load_model('mnist_model.h5')

class RecognizeApp(tk.Toplevel):
    def __init__(self, master=None):
        super().__init__(master)
        self.setup_ui()
        self.setup_bindings()
        self.init_canvas_image()

    def setup_ui(self):
        """设置用户界面组件。"""
        self.title('MNIST手写数字识别')
        self.geometry('400x450')

        self.canvas = tk.Canvas(self, bg='white', width=280, height=280)
        self.canvas.pack(pady=20)

        self.button_recognize = tk.Button(self, text='识别', command=self.recognize)
        self.button_recognize.pack()

        self.button_clear = tk.Button(self, text='清空', command=self.clear)
        self.button_clear.pack()

        self.button_load = tk.Button(self, text='加载图片', command=self.load_image)
        self.button_load.pack()

        self.label_result = tk.Label(self, text='', font=('Helvetica', 20))
        self.label_result.pack(pady=20)

    def setup_bindings(self):
        """设置事件绑定。"""
        self.canvas.bind('<B1-Motion>', self.paint)

    def init_canvas_image(self):
        """初始化画布的基础图像和绘制工具。"""
        self.image = Image.new('L', (280, 280), 255)
        self.draw = ImageDraw.Draw(self.image)

    def paint(self, event):
        """处理鼠标拖动事件，在画布上绘制圆点。"""
        x1, y1 = (event.x - 5), (event.y - 5)
        x2, y2 = (event.x + 5), (event.y + 5)
        self.canvas.create_oval(x1, y1, x2, y2, fill='black', width=10)
        self.draw.ellipse([x1, y1, x2, y2], fill='black')

    def recognize(self):
        """识别画布上的手写数字并在标签中显示结果。"""
        img = self.preprocess_image(self.image)
        predictions = model.predict(img)
        digit, confidence = self.get_prediction_result(predictions)

        # 打印调试信息
        print(f'Debug: Predictions: {predictions}, Digit: {digit}, Confidence: {confidence}')

        # 更新结果标签
        self.label_result.config(text=f'识别结果: {digit}, 置信度: {confidence:.2f}')
        self.update_idletasks()

    def preprocess_image(self, image):
        """预处理图像以供模型预测。"""
        img = image.resize((28, 28))
        img = img.filter(ImageFilter.MedianFilter())  # 去噪
        img = ImageOps.invert(img)  # 反转颜色
        img = ImageOps.autocontrast(img)  # 增强对比度
        img = np.array(img) / 255.0
        img = img.reshape(1, 28, 28, 1)
        return img

    def get_prediction_result(self, predictions):
        """从模型预测结果中提取数字和置信度。"""
        digit = np.argmax(predictions)
        confidence = np.max(predictions)
        return digit, confidence

    def clear(self):
        """清空画布和结果标签。"""
        self.canvas.delete('all')
        self.init_canvas_image()
        self.label_result.config(text='')

    def load_image(self):
        """从文件系统加载图像并进行识别。"""
        file_path = filedialog.askopenfilename()
        if not file_path:
            return
        try:
            loaded_image = Image.open(file_path).convert('L')
            preprocessed_image = self.preprocess_image(loaded_image)
            predictions = model.predict(preprocessed_image)
            digit, confidence = self.get_prediction_result(predictions)

            # 打印调试信息
            print(f'Debug: Predictions: {predictions}, Digit: {digit}, Confidence: {confidence}')

            # 更新结果标签
            self.label_result.config(text=f'识别结果: {digit}, 置信度: {confidence:.2f}')
            self.update_idletasks()
        except Exception as e:
            messagebox.showerror("加载错误", f"无法加载图像: {e}")

if __name__ == "__main__":
    root = tk.Tk()
    app = RecognizeApp(master=root)
    root.mainloop()